/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.client;

import java.rmi.NoSuchObjectException;
import java.rmi.RemoteException;
import java.rmi.server.UnicastRemoteObject;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.UUID;

import org.slf4j.Logger;
import org.sss.common.JobDoneListener;
import org.sss.common.RemoteStatus;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;

/**
 * Job-completion waiter.
 * 
 * NOTE: This class is used by only this module.
 */
class JobWaiter<T extends RemoteStatus> {
  private static Logger logger = SssClient.getLogger(JobWaiter.class);
  private final Object lock = new Object();
  private final Map<StorageNode, Listener<T>> listeners = new HashMap<StorageNode, Listener<T>>();
  private final ClusterManager clusterManager;
  private boolean alreadyCanceled = false;
  private final UUID jobid;
  private Throwable keepAliveError = null;

  /**
   * Constructor of  this class.
   *
   * @param client SssClient.
   * @param jobid  Identifier of job.
   */
  public JobWaiter(SssClient client, UUID jobid) throws SssException {
    this.clusterManager = client.getClusterManager();
    this.jobid = jobid;
    client.registerJobWaiter(this);
    clusterManager.register(this);
  }

  /**
   * Waits completion of job on all remote hosts.
   *
   * @return map of remote host and job-result.
   */
  public Map<StorageNode, T> await() throws SssException, InterruptedException {
    try {
      synchronized (lock) {
        Map<StorageNode, T> result = new HashMap<StorageNode, T>();

        while (check(result)) {
          lock.wait();
        }
        return result;
      }
    } finally {
      clusterManager.unregister(this);
    }
  }

  /**
   * Creates new job-completion listener.
   */
  public synchronized Listener<T> newListener(StorageNode node) throws RemoteException {
    Listener<T> l = new Listener<T>(lock);
    listeners.put(node, l);

    return l;
  }

  /**
   * Dispose this.
   */
  public void dispose() {
    Iterator<Map.Entry<StorageNode, Listener<T>>> it = listeners.entrySet().iterator();
    while (it.hasNext()) {
      it.next().getValue().dispose();
      it.remove();
    }
  }

  @Override
  public void finalize() {
    dispose();
  }

  /**
   *.
   * NOTE: Module level private.
   */
  void setKeepAliveError(Throwable e) {
    synchronized (lock) {
      keepAliveError = e;
      lock.notifyAll();
    }
  }

  private boolean check(Map<StorageNode, T> result) throws SssException {
    synchronized (lock) {
      if (keepAliveError != null) {
        throw new SssException("Keep alive error.", keepAliveError);
      }
      boolean reqCancel = false;
      Iterator<Map.Entry<StorageNode, Listener<T>>> it = listeners.entrySet().iterator();
      while (it.hasNext()) {
        Map.Entry<StorageNode, Listener<T>> e = it.next();
        if (e.getValue().status != null) {
          it.remove();
          e.getValue().dispose();
          T status = e.getValue().status;
          result.put(e.getKey(), status);

          if (status.exception != null) {
            reqCancel = true;
          }
        }
      }
      if (reqCancel && !alreadyCanceled) {
        alreadyCanceled = true;
        cancel();
      }
      return !listeners.isEmpty();
    }
  }
  
  /**
   * Cancel target job.
   */
  public void cancel() {
    logger.debug("cancel tasks.");
    synchronized (lock) {
      for (Map.Entry<StorageNode, Listener<T>> e: listeners.entrySet()) {
        if (e.getValue().status == null) {
          try {
            logger.info("Cancel {}.", e.getKey());
            ClusterManager.Server server = clusterManager.get(e.getKey());
            try {
              server.get().cancel(jobid);
            } finally {
              server.unlock();
            }
          } catch (Throwable ex) {
            // TODO: Throw outside?
            logger.error("Can't cancel.", ex);
          }
        }
      }
    }
  }

  /**
   * Job-completion listener
   *
   * NOTE: This clas is used in internal. Users don't need to touch this class.
   */
  @SuppressWarnings("serial")
  public static class Listener<T extends RemoteStatus> extends UnicastRemoteObject implements JobDoneListener<T> {
    private final Object lock;
    private T status = null;

    private Listener(Object lock) throws RemoteException {
      this.lock = lock;
    }

    /**
     * Callback method.
     *
     * This is called when job is completed.
     */
    @Override
    public void jobDone(T status) throws RemoteException {
      synchronized (lock) {
        assert status != null;
        this.status = status;
        lock.notifyAll();
      }
    }

    private void dispose() {
      try {
        UnicastRemoteObject.unexportObject(this, true);
      } catch (NoSuchObjectException e) {
        logger.warn("", e);
      }
    }
  }
}
